package com.lkx.kafka.monitor.common;

import com.lkx.kafka.monitor.dto.ThroughputMetrics;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

@Component
@Slf4j
public class MemoryDataManager {
    private final Map<String, List<ThroughputMetrics>> memoryStore = new ConcurrentHashMap<>();

    /**
     * 添加数据到内存
     */
    public void addData(ThroughputMetrics metrics) {
        String topicName = metrics.getTopicName() != null ? metrics.getTopicName() : "ALL";
        memoryStore.computeIfAbsent(topicName, k -> new ArrayList<>()).add(metrics);
    }

    /**
     * 获取指定主题的数据
     */
    public List<ThroughputMetrics> getData(String topicName, Long startTime, Long endTime) {
        String key = topicName != null ? topicName : "ALL";
        List<ThroughputMetrics> dataList = memoryStore.getOrDefault(key, new ArrayList<>());

        return dataList.stream()
                .filter(data -> {
                    boolean afterStart = startTime == null || data.getTimestamp() >= startTime;
                    boolean beforeEnd = endTime == null || data.getTimestamp() <= endTime;
                    return afterStart && beforeEnd;
                })
                .sorted(Comparator.comparing(ThroughputMetrics::getTimestamp))
                .collect(Collectors.toList());
    }

    /**
     * 获取所有主题的数据
     */
    public Map<String, List<ThroughputMetrics>> getAllData() {
        return new HashMap<>(memoryStore);
    }

    /**
     * 设置数据（用于从文件加载）
     */
    public void setData(Map<String, List<ThroughputMetrics>> data) {
        memoryStore.clear();
        memoryStore.putAll(data);
    }

    /**
     * 清空数据
     */
    public void clear() {
        memoryStore.clear();
    }
}
